import torch
from transformers import BertTokenizer, BertForSequenceClassification
import pandas as pd
import argparse
import os
import logging

# 设置日志
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%Y/%m/%d %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)

class SentimentPredictor:
    def __init__(self, model_path):
        """
        初始化情感预测器
        
        参数:
            model_path: 微调后的模型路径
        """
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        logger.info(f"使用设备: {self.device}")
        
        # 加载模型和tokenizer
        self.tokenizer = BertTokenizer.from_pretrained(model_path)
        self.model = BertForSequenceClassification.from_pretrained(model_path)
        self.model.to(self.device)
        self.model.eval()
        
        logger.info(f"模型加载完成: {model_path}")
        
        # 情感标签映射
        self.id2label = {0: '负面', 1: '中性', 2: '正面'}
    
    def predict(self, text, max_length=128):
        """
        预测文本的情感
        
        参数:
            text: 输入文本
            max_length: 最大序列长度
            
        返回:
            sentiment: 情感标签
            confidence: 置信度
        """
        # 对文本进行编码
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=max_length,
            return_token_type_ids=True,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt'
        )
        
        # 将编码放到指定设备上
        input_ids = encoding['input_ids'].to(self.device)
        attention_mask = encoding['attention_mask'].to(self.device)
        token_type_ids = encoding['token_type_ids'].to(self.device)
        
        # 预测
        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                token_type_ids=token_type_ids,

            )
        
        # 获取预测结果
        logits = outputs.logits
        probabilities = torch.nn.functional.softmax(logits, dim=1)
        confidence, predicted_class = torch.max(probabilities, dim=1)
        
        sentiment = self.id2label[predicted_class.item()]
        confidence = confidence.item()
        
        return sentiment, confidence

def predict_file(predictor, input_path, output_path=None):
    """
    预测CSV文件中的文本情感
    
    参数:
        predictor: 情感预测器实例
        input_path: 输入CSV文件路径
        output_path: 输出CSV文件路径
    """
    logger.info(f"开始处理文件: {input_path}")
    
    # 读取CSV文件
    df = pd.read_csv(input_path)
    
    # 确保post_title列存在
    if 'post_title' not in df.columns:
        logger.error("CSV文件中缺少'post_title'列")
        return
    
    logger.info(f"需要预测的行数: {len(df)}")
    
    # 添加新列用于存储预测结果
    df['predicted_sentiment'] = None
    df['confidence'] = None
    
    # 逐行预测
    for idx, row in df.iterrows():
        text = str(row['post_title'])
        sentiment, confidence = predictor.predict(text)
        
        df.at[idx, 'predicted_sentiment'] = sentiment
        df.at[idx, 'confidence'] = confidence
        
        if idx % 100 == 0:
            logger.info(f"已处理 {idx}/{len(df)} 行")
    
    # 如果没有指定输出文件，生成默认输出路径
    if output_path is None:
        base_name = os.path.basename(input_path)
        name_without_ext = os.path.splitext(base_name)[0]
        output_path = f"{name_without_ext}_predicted.csv"
    
    # 保存结果
    df.to_csv(output_path, index=False, encoding='utf-8-sig')
    logger.info(f"预测完成，结果已保存到: {output_path}")
    
    # 显示情感分布
    sentiment_counts = df['predicted_sentiment'].value_counts()
    logger.info(f"预测情感分布:\n{sentiment_counts}")
    
    # 如果有真实情感标签，计算一致性
    if 'sentiment' in df.columns:
        match_count = (df['predicted_sentiment'] == df['sentiment']).sum()
        match_rate = match_count / len(df)
        logger.info(f"与真实标签一致数: {match_count}/{len(df)} ({match_rate:.2%})")
    
    return df

def predict_text(predictor, text):
    """
    预测单个文本的情感
    
    参数:
        predictor: 情感预测器实例
        text: 输入文本
    """
    sentiment, confidence = predictor.predict(text)
    logger.info(f"文本: {text}")
    logger.info(f"预测情感: {sentiment} (置信度: {confidence:.4f})")
    return sentiment, confidence

def main():
    parser = argparse.ArgumentParser(description="使用微调后的FinBERT模型预测情感")
    
    parser.add_argument("--model_path", type=str, required=True,
                      help="微调后的模型路径")
    parser.add_argument("--input_file", type=str, default=None,
                      help="输入CSV文件路径")
    parser.add_argument("--output_file", type=str, default=None,
                      help="输出CSV文件路径")
    parser.add_argument("--text", type=str, default=None,
                      help="要预测的文本")
    
    args = parser.parse_args()
    
    # 初始化预测器
    predictor = SentimentPredictor(args.model_path)
    
    # 根据输入参数决定预测文件还是单个文本
    if args.input_file:
        predict_file(predictor, args.input_file, args.output_file)
    elif args.text:
        predict_text(predictor, args.text)
    else:
        logger.error("请提供输入文件或文本进行预测")

if __name__ == "__main__":
    main() 